"""
detect.py
"""
from executor.predict import predict_single_img
from server.api_params_def import *
import base64
import numpy as np
import cv2
from io import BytesIO
from PIL import Image
import traceback

from commons.logger import logger


def base64_to_image(base64_str: str) -> np.ndarray:
    """
    将base64字符串转换为numpy数组格式的图片

    @param base64_str: base64编码的图片字符串
    @return: numpy.ndarray 格式的图片
    """
    # 如果base64字符串包含前缀（如 'data:image/jpeg;base64,'），需要去除
    if ',' in base64_str:
        base64_str = base64_str.split(',')[1]

    try:
        # base64解码
        img_data = base64.b64decode(base64_str)
        # 转换为numpy数组
        img_array = np.frombuffer(img_data, np.uint8)
        # 解码图片
        img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
        # OpenCV使用BGR格式，转换为RGB格式
        # return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        return img
    except Exception as e:
        raise Exception(f"图片转换失败: {e}")


def process_detect(input_params: DetectInputParams, config: ConfigInfo) -> DetectOutputParams:
    """
    处理诊断请求

    Args:
        input_params: 输入参数
        config: 配置信息

    Returns:
        DetectOutputParams: 诊断输出结果
    """
    try:
        start_time = datetime.now()

        # 转换图片格式
        np_image = base64_to_image(input_params.image_base64)

        # 目标检测
        detect_vals, predict_img = predict_single_img(np_image,
                                                      conf_threshold=config.conf_threshold,
                                                      plot_result=config.plot_result)
        logger.info(f"detected vals len: {len(detect_vals)}")

        # 处理检测结果 - 如果有多个检测结果，选择置信度最高的一个
        score = 0.0
        detect_code = ""
        detect_vertices = np.array([0, 0, 0, 0])  # 默认值

        if detect_vals:
            # 获取置信度最高的检测结果
            best_detection = max(detect_vals, key=lambda x: x['conf'])
            score = best_detection['conf']
            detect_code = best_detection['cls']  # 类别作为检测代码
            detect_vertices = np.array(best_detection['box'])  # 边界框坐标

        # 处理结果图像
        if config.plot_result:
            # predict_img 中已经包含了绘制的检测框，因为predict_single_img函数在plot_result=True时会绘制
            # 转base64
            pil_image = Image.fromarray(predict_img)
            buffer = BytesIO()
            pil_image.save(buffer, format="PNG")
            result_base64 = "data:image/jpeg;base64," + base64.b64encode(buffer.getvalue()).decode()
        else:
            result_base64 = ""

        predict_ms = int((datetime.now() - start_time).total_seconds() * 1000)

        # 组装输出结果，现在包含了result_base64字段
        return DetectOutputParams(
            error_code=0,
            error_msg="",
            timestamp=datetime.now(),
            predict_ms=predict_ms,
            score=score,
            detect_code=detect_code,
            detect_vertices=detect_vertices,
            result_base64=result_base64
        )

    except Exception as e:
        # 获取异常的详细信息
        tb = traceback.format_exc()  # 获取完整的 traceback 信息
        error_msg = f"{str(e)}\nTraceback:\n{tb}"
        # 发生异常时返回错误信息
        return DetectOutputParams(
            error_code=1,
            error_msg=error_msg,
            timestamp=datetime.now(),
            predict_ms=0,
            score=0.0,
            detect_code="",
            detect_vertices=np.array([0, 0, 0, 0]),
            result_base64=""  # 错误情况下返回空字符串
        )